import numpy

from physics.vector.array import _array, Vector3Array




class Rotation3Array(object):
    '''
    NumPy-like array of 3D rotation matrices.

    Parameters
    ----------
    w : Length-N sequence of 3x3 matrices
    '''

    def __init__(self, *args, **kwargs):        
       if 'dtype' in kwargs:
           self.dtype = kwargs['dtype']
       else:
           self.dtype = numpy.float64

       if len(args) == 1:
           if numpy.isscalar(args[0]):
               # Identity array of length N
               self.r = numpy.array([numpy.identity(3)]*args[0])

           else:
               r = numpy.array(args[0])

               # Shape should be (N, 3, 3)
               if r.shape[1:] != (3,3):
                   msg = "Calling Rotation3Array.__init__() with single argument expects that argument"
                   msg += " to be of shape (N,3,3); shape given is {0}".format(r.shape)
                   raise ValueError(msg)

               self.r = r

       elif len(args) == 3:           

           if r.shape[1:] != (3,3):
               msg = "Calling Rotation3Array.__init__() with single argument expects that argument"
               msg += " to be of shape (N,3,3); shape given is {0}".format(r.shape)
           

    # This function is for the Vector3Array
    def rotate(self, angle, axis):
        '''
        Rotates each vector *in place* by angle `angle' around an arbitrary axis `axis'.

        Parameters:
        --------------------
        angle: float
            The angle of rotation about `axis'
        axis: 3-tuple, sequence of 3-tuples of length N, or Vector3Array of length N
            The axis (or array of axes) to rotate each vector about.
        '''   
        
        if type(axis) == type(self):
            if axis.shape != self.shape:
                msg = "Axis shape {0} doesn't match data shape {1}".format(axis.shape, data.shape)
                raise ValueError(msg)
            
            return _array.v3_rotate(self.x, self.y, self.z, 
                                    axis.x, axis.y, axis.z, 
                                    angle)

        else:
            axis = numpy.array(axis)

            if len(axis.shape) == 1:
                if axis.shape[0] != 3:
                    msg = "Argument 2 is a 1D sequence of length {0}; must be length 3"
                    raise ValueError(msg.format(axis.shape[0]))
                
                return _array.v3_rotate_single_axis(self.x, self.y, self.z, 
                                                    axis[0], axis[1], axis[2],
                                                    angle)
            elif axis.shape != self.shape:
                msg = "Argument 2 has shape {0}; must have same shape as data ({1})"
                raise ValueError(msg.format(axis.shape, self.shape))
            
            else:
                tmp = Vector3Array(axis)
                return _array.v3_rotate(self.x, self.y, self.z, axis.x, 
                                        axis.x, axis.y, axis.z,
                                        angle)
